# This Python file uses the following encoding: utf-8

import h5py as h5
from template import *
import math



class TriMesh2D:
    def __init__(self):
        pass

    def read_cgns_mesh(self, file_name):
        mesh_file = h5.File(file_name, "r")
	
        list_keys = list(mesh_file.keys())
        print("list_keys: ", list_keys)
        print("hdf version: ", mesh_file[" hdf5version"][:])
        print("CGNSLibraryVersion version: ", mesh_file["CGNSLibraryVersion"])
	
	
        mesh_1 = mesh_file["Mesh_1"]
        #mesh_1 = mesh_file["t1"]
        print("mesh_1 keys: ", list(mesh_1.keys()))

        smesh_mesh = mesh_1["SMESH_Mesh"]

        edge_element = list(smesh_mesh.keys())[3]
        tri_element  = list(smesh_mesh.keys())[4]

        self.vertex_x = smesh_mesh["GridCoordinates"]["CoordinateX"][" data"][:]
        self.vertex_y = smesh_mesh["GridCoordinates"]["CoordinateY"][" data"][:]
        self.vertex_z = smesh_mesh["GridCoordinates"]["CoordinateZ"][" data"][:]

        self.bar_elements = smesh_mesh[edge_element]["ElementConnectivity"][" data"][:]
        self.bar_elements = self.bar_elements.reshape((int(self.bar_elements.shape[0]/2), 2)) - 1

        self.tri_elements = smesh_mesh[tri_element]["ElementConnectivity"][" data"][:]
        self.tri_elements = self.tri_elements.reshape((int(self.tri_elements.shape[0]/3), 3)) - 1

    def read_pmsh_mesh(self, file_name):
        mesh_file = h5.File(file_name, "r")
        mesh2d = mesh_file["mesh2d"]

        vertexes = mesh2d["vertexes"][:]

        self.vertex_x = vertexes[:, 0]
        self.vertex_y = vertexes[:, 1]
	
        self.tri_elements = mesh2d["tri_elements"][:]
	
        self.bar_elements = mesh2d["bar_elements"][:]

        edge_group = mesh_file["edge"]
        self.edge_name_list = list(edge_group.keys())
        self.edge_list = {}
        for edge_name in self.edge_name_list:
            self.edge_list[edge_name] = edge_group[edge_name][:]
            #print(edge_name, " shape: ", self.edge_list[edge_name].shape)
            #print(edge_name, " bar_list: ", self.edge_list[edge_name])
	

        #cal bar_x, bar_y, bar_distance
        self.bar_x = np.zeros(self.bar_elements.shape[0])
        self.bar_y = np.zeros(self.bar_elements.shape[0])
        self.bar_distance = np.zeros(self.bar_elements.shape[0])
        for i in range(0, self.bar_elements.shape[0]):
            bar0 = self.bar_elements[i, 0]
            bar1 = self.bar_elements[i, 1]
            x0 = self.vertex_x[bar0]
            y0 = self.vertex_y[bar0]
            x1 = self.vertex_x[bar1]
            y1 = self.vertex_y[bar1]
            self.bar_x[i] = 0.5 * (x0 + x1)
            self.bar_y[i] = 0.5 * (y0 + y1)

            bar_length = math.sqrt( (x1-x0) * (x1-x0) + (y1-y0) * (y1-y0) )
            if i == 0:
                self.bar_distance[i] = 0.5 * bar_length
            else:
                self.bar_distance[i] = self.bar_distance[i-1] + 0.5 * (bar_length_pre + bar_length)
            bar_length_pre = bar_length

        mesh_file.close()


    def plot_mesh(self, index0 = 0, index1 = 0):

        fig=plt.figure(figsize=(10,8))
        fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.4)

        #------ potential ---------------
        ax0 = fig.add_subplot(1,1,1)

        ax0.triplot(self.vertex_x, self.vertex_y, self.tri_elements, lw=1.0)
        #ax0.triplot(self.vertex_x, self.vertex_y, self.tri_elements[0:1, :], lw=1.0)

        if index1 > index0:
            edge_points = np.zeros((index1 - index0, 2))
            for i in range(index0, index1):
                edge_points[i-index0, 0] = self.vertex_x[self.bar_elements[i, 0]]
                edge_points[i-index0, 1] = self.vertex_y[self.bar_elements[i, 0]]
            ax0.scatter(edge_points[:, 0], edge_points[:, 1])
        ax0.set_aspect('equal')


        fig.savefig("mesh.png", dpi = 300)
        plt.show()





'''
if__name__ == "__main__":
    mesh = TriMesh2D ()
    mesh.read_cgns_mesh("mesh100.cgns")
    mesh.plot_mesh()
'''
